import gymnasium as gym
from stable_baselines3 import PPO
import torch

# 确认设备
device = "cuda" if torch.cuda.is_available() else "cpu"
print("Using device:", device)

# 创建环境（带图形渲染）
env = gym.make("LunarLander-v3", render_mode="human")


# 创建 PPO 模型
model = PPO(
    "MlpPolicy",
    env,
    verbose=1,
    device=device,
    tensorboard_log="./ppo_log/"
)

# 训练 5 万步（大约几分钟）
print("开始训练...")
model.learn(total_timesteps=50000)
print("训练结束 ✅")

# 保存模型
model.save("ppo_lunarlander")

# 测试模型
obs, _ = env.reset()
for step in range(2000):
    action, _ = model.predict(obs, deterministic=True)
    obs, reward, terminated, truncated, _ = env.step(action)
    if terminated or truncated:
        obs, _ = env.reset()

env.close()

